
from common.utils import *
import os
import cv2
import xml.dom.minidom
from xml.dom.minidom import parse
import xml.etree.ElementTree as ET
import sys
import os


def read_labelimg_xml(img_path, readimg=False):
    # 载入数据（root得到根节点，打印得到的是内存地址）
    tree = ET.parse(img_path)
    root = tree.getroot()
    out = {}
    # 建立一个字典
    for annotation in root.iter('annotation'):
        # 得到object中的信息
        for ob in annotation.iter('object'):
            name = ''
            # 得到标注框名称信息
            for namex in ob.iter('name'):
                name = namex.text

            if name not in out:
                out[name] = []

            x1, y1, x2, y2 = 0, 0, 0, 0
            rects = []
            # 得到标注框坐标位置信息
            for bndbox in ob.iter('bndbox'):
                for xx in bndbox.iter('xmin'):
                    x1 = float(xx.text)

                for xx in bndbox.iter('ymin'):
                    y1 = float(xx.text)

                for xx in bndbox.iter('xmax'):
                    x2 = float(xx.text)

                for xx in bndbox.iter('ymax'):
                    y2 = float(xx.text)

                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                r = (x1, y1, x2, y2)
                rects.append(r)

            out[name] += rects

    return out


def draw_labelimg_xml(img, out):
    # 绘制矩形框
    for aa in out:
        objs = aa['objs']
        for obj in objs:
            xmin, ymin = 0, 0
            for r in obj['rects']:
                xmin, ymin = r[0], r[1]
                cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), (0, 255, 0), 2)

            font = cv2.FONT_HERSHEY_DUPLEX

            name1 = obj['name']
            # 字符串绘制函数
            cv2.putText(img, name1, (xmin, ymin),
                        font, 1.2, (255, 255, 255), 2)
            #savename = dir[:-4] + ".jpg"

    return img


def draw_labelimg_xml_neg(img, out):
    # 绘制矩形框
    for aa in out:
        objs = aa['objs']
        for obj in objs:
            xmin, ymin = 0, 0
            for r in obj['rects']:
                xmin, ymin = r[0], r[1]
                cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), (0, 0, 0), -1)

    return img


def rect_norm(r, h, w):
    rh = r[3]-r[1]
    rw = r[2]-r[0]
    r = list(r)
    if rh*w < rw*h:
        rh2 = int(rw*h/w)
        r[1] -= (rh2-rh)//2
        r[3] = r[1]+rh2
    else:
        rw2 = int(rh*w/h)
        r[0] -= (rw2-rw)//2
        r[2] = r[0]+rw2

    return tuple(r)


def vadd(a, b):
    assert len(a) == len(b)
    c = []
    for i in range(len(a)):
        c.append((a[i]+b[i])//2)

    return tuple(c)


def clip_labelimg_xml(img, out, fn1):
    cnt = 0
    mkdir('ppp')
    # GaussianNoise(0, 5), Leaf(),
    ops = [GaussianBlur((7, 7)), Vignetting(), LensDistortion()]
    for aa in out:
        objs = aa['objs']
        for obj in objs:
            name1 = obj['name']
            for r in obj['rects']:
                if '_' in name1:
                    h = r[3]-r[1]
                    w = (r[2]-r[0])//len(name1)
                    if h < w*1.1:
                        h1 = int(w*1.5)
                        r1 = r[3] - h1
                        r = (r[0], r1, r[2], r[3])
                #cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), (0, 255, 0), 2)
                #img1 = imrect(img, r)
                r1 = r

            r2 = cut_rect(r1, 1, len(name1))
            if len(name1) >= 2 and True:
                name1 = name1+'~'
                r2.append(vadd(r2[0], r2[1]))

            #s = get_time_stamp()
            for i in range(len(name1)):
                r = r2[i]
                r = rect_norm(r, 3, 2)
                im = imrect(img, r)
                c = name1[i]
                mkdir('pic/%s' % (c))
                #cv2.imwrite('pic/%s/%s_%d_%d.jpg' % (c, fn1, cnt, i), im)
                cv2.imwrite('pic/%s/%s.jpg' % (c, fn1), im)

                mkdir('ppp/%s' % (c))
                imgs = imarg(img, r, 10, ops, 0.05, 0.05, 5)
                for j in range(len(imgs)):
                    im = imgs[j].astype(np.uint8)
                    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
                    im = cv2.resize(im, (20, 30))
                    cv2.imwrite('ppp/%s/%s_%d.jpg' % (c, fn1, j), im)

            cnt += 1

            # 字符串绘制函数
            #cv2.imwrite(save_dir+savename, img)

    return img


def read_all_labelimg_xml(img_dir, readimg=False):
    dirs = os.listdir(img_dir)
    outs = []
    for dir in dirs:
        pre_dir = dir[-3:]

        if not pre_dir == 'xml':
            continue

        img_path = img_dir+dir
        print(img_path)
        out = read_labelimg_xml(img_path, readimg)
        outs.append(out)

    return outs


def test_read_labelimg_xml():
    # 读入图片路径（以/结束）
    # os.chdir('D:/data/通力电梯七段管OCR/111/test1')
    os.chdir('D:/data/通力电梯七段管OCR/tonlipic/image_biaozhu_test')
    mkdir('out')
    mkdir('neg')
    readimg = True
    names = os.listdir('./')
    for i in range(len(names)):
        fn = names[i]

        if '.jpg' not in fn:
            continue

        xml_path = fn.replace('.jpg', '.xml')
        if not os.path.exists(xml_path):
            continue

        fn1 = fn.replace('.bmp', '').replace('.jpg', '')
        print(fn)
        out = read_labelimg_xml(xml_path, readimg)
        for aa in out:
            img = aa['img']

        #clip_labelimg_xml(img, out, fn1)
        #img1 = draw_labelimg_xml(img, out)
        #cv2.imwrite('out/%s_img.jpg' % fn1, draw_labelimg_xml(img, out))
        cv2.imwrite('neg/%s_img.jpg' % fn1, draw_labelimg_xml_neg(img, out))


if __name__ == "__main__":
    test_read_labelimg_xml()
